#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <opencv2/opencv.hpp>
#include "rknn_api.h" 
#include <algorithm>
#include <cmath>

// 检测框结构体
struct Bbox {
    float x1, y1;  // 左上角坐标
    float x2, y2;  // 右下角坐标
    float conf;    // 置信度
    int cls;       // 类别索引
};

class RKNNDetector 
{
    public:
        // 构造函数：初始化参数
        RKNNDetector(const std::string& model_path, 
                     const int input_width = 320, 
                     const int input_height = 320, 
                     const float conf_thresh = 0.5f, 
                     const float nms_thresh = 0.45f, 
                     int nc = 1);
    
        // 析构函数：释放资源
        ~RKNNDetector();
    
        // 初始化：加载模型并准备推理环境
        bool init();
    
        // 处理单帧图像：返回检测结果
        std::vector<Bbox> process(cv::Mat& frame);
    
        // 在图像上绘制检测结果
        void draw(cv::Mat& frame, const std::vector<Bbox>& boxes);
    
        // 手动释放资源
        void release();
    
        // 判断是否初始化成功
        bool isInited() const;
    
    private:
        // 模型配置参数
        std::string model_path_;
        int INPUT_WIDTH_;
        int INPUT_HEIGHT_;
        float CONF_THRESH_;
        float NMS_THRESH_;
        int nc_;  // 类别数
        std::vector<std::string> CLASSES_;  // 类别名称
    
        // YOLOv5锚框和下采样倍数（与训练一致） 320
        const std::vector<std::vector<std::pair<float, float>>> ANCHORS_ = {
            {{10, 13}, {16, 30}, {33, 23}},
            {{30, 61}, {62, 45}, {59, 119}},
            {{116, 90}, {156, 198}, {373, 326}}
        };
        // 修改后（640×640，按2倍缩放）
        // const std::vector<std::vector<std::pair<float, float>>> ANCHORS_ = {
        //     {{20, 26}, {32, 60}, {66, 46}},    // 10×2=20，13×2=26，以此类推
        //     {{60, 122}, {124, 90}, {118, 238}},
        //     {{232, 180}, {312, 396}, {746, 652}}
        // };

        const std::vector<int> STRIDES_ = {8, 16, 32};
    
        // RKNN相关变量
        rknn_context ctx_;
        uint32_t n_input_;
        uint32_t n_output_;
        std::vector<rknn_tensor_attr> output_attrs_;
        std::vector<int> output_sizes_precomputed_;
        std::vector<rknn_input> inputs_;
        std::vector<rknn_output> outputs_;
        bool is_inited_;
        float input_scale_; // 输入张量的量化缩放因子
        int32_t input_zero_point_;  // 输入张量的量化零点(偏移量)
    
        // 图像预处理：缩放、填充、格式转换
        cv::Mat preprocess(cv::Mat& frame, float& scale, float& pad_w, float& pad_h);
    
        // 解码单个锚框的预测结果
        void decode_box(float* feat_ptr, const std::pair<float, float>& anchor,
                       float scale, float pad_w, float pad_h,
                       int grid_w, int grid_h, int stride,
                       std::vector<Bbox>& boxes);
    
        // 后处理：解析多尺度特征图并生成检测框
        std::vector<Bbox> postprocess(std::vector<float*>& outputs,
                                     std::vector<int>& output_sizes,
                                     float scale, float pad_w, float pad_h,
                                     const cv::Mat& frame);
    
        // 计算IOU
        float calculate_iou(const Bbox& box1, const Bbox& box2);
    
        // 非极大值抑制（NMS）
        std::vector<Bbox> apply_nms(std::vector<Bbox>& boxes, float iou_threshold);
        // 辅助函数：判断框a是否包含框b
        bool contains(const Bbox& a, const Bbox& b);
};
    